# Author: WY
# Date: 2023/5/17 16:28

import torch.nn as nn
from torch.nn.utils import weight_norm


'''
# 超参数
scopes = 50
lr = 0.0002 # 0.0003
split = 0.9 # 0.9
weight_decay = 5e-4
feature_dim = 13
batch_size = 40 # 120
output_size = 20
input_size = batch_size - output_size
channels = [50, 50, 50] # [50, 50, 50]
kernel_size = 3
stride = 1
dropout = 0.2
'''


class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size

    def forward(self, x):
        # return x[:, :, :-self.chomp_size].contiguous()
        return x[:, :-self.chomp_size].contiguous()


class TemporalBlock(nn.Module):
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        """
        相当于一个Residual block

        :param n_inputs: int, 输入通道数
        :param n_outputs: int, 输出通道数
        :param kernel_size: int, 卷积核尺寸
        :param stride: int, 步长，一般为1
        :param dilation: int, 膨胀系数
        :param padding: int, 填充系数
        :param dropout: float, dropout比率
        """
        super(TemporalBlock, self).__init__()
        # weight_norm: 权重归一化
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation))
        # self.conv1 = weight_norm(nn.Conv2d(n_inputs, n_outputs, (1, kernel_size), stride=stride, padding=padding, dilation=dilation))
        self.softmax1 = nn.Softmax(dim=0)
        # 经过conv1，输出的size其实是(Batch, input_channel, seq_len + padding)
        self.chomp1 = Chomp1d(padding)  # 裁剪掉多出来的padding部分，维持输出时间步为seq_len
        self.relu1 = nn.ReLU(inplace=True)
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation))
        # self.conv2 = weight_norm(nn.Conv2d(n_outputs, n_outputs, (1, kernel_size), stride=stride, padding=padding, dilation=dilation))
        self.softmax2 = nn.Softmax(dim=0)
        self.chomp2 = Chomp1d(padding)  # 裁剪掉多出来的padding部分，维持输出时间步为seq_len
        self.relu2 = nn.ReLU(inplace=True)
        self.dropout2 = nn.Dropout(dropout)

        self.net = nn.Sequential(
            self.conv1, self.softmax1, self.chomp1, self.relu1, self.dropout1,
            self.conv2, self.softmax2, self.chomp2, self.relu2, self.dropout2
        )

        # 卷积核为1x1的卷积过程,在输入TemporalBlock时数据的通道数与输出时的通道数不等时使用,计算残差
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.relu = nn.ReLU(inplace=True)
        self.init_weights()

    def init_weights(self):
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        if self.downsample is not None:
            self.downsample.weight.data.normal_(0, 0.01)

    def forward(self, x):
        # out = self.conv1(x)
        # out = self.chomp1(out)
        # return self.relu(out)

        out = self.net(x)
        res = x if self.downsample is None else self.downsample(x)
        return self.relu(out + res)


class TemporalConvNet(nn.Module):
    def __init__(self,  input_size, channels, kernel_size, stride, dropout):
        """
        TCN，目前paper给出的TCN结构很好的支持每个时刻为一个数的情况，即sequence结构，
        对于每个时刻为一个向量这种一维结构，勉强可以把向量拆成若干该时刻的输入通道，
        对于每个时刻为一个矩阵或更高维图像的情况，就不太好办。

        :param input_size: int， 输入通道数
        :param channels: list，每层的hidden_channel数，例如[25,25,25,25]表示有4个隐层，每层hidden_channel数为25
        :param kernel_size: int, 卷积核尺寸
        :param stride: int, 卷积步长
        :param dropout: float, drop_out比率
        """
        super(TemporalConvNet, self).__init__()
        layers = []
        num_levels = len(channels)
        for i in range(num_levels):
            dilation_size = 2 ** i  # 膨胀系数：1，2，4，8……
            in_channels = input_size if i == 0 else channels[i - 1]  # 确定每一层的输入通道数
            out_channels = channels[i]  # 确定每一层的输出通道数
            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation_size, padding=(kernel_size - 1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        """
        输入x的结构不同于RNN，一般RNN的size为(Batch, seq_len, channels)或者(seq_len, Batch, channels)，
        这里把seq_len放在channels后面，把所有时间步的数据拼起来，当做Conv1d的输入尺寸，实现卷积跨时间步的操作

        :param x: size of (Batch, input_channel, seq_len)
        :return: size of (Batch, output_channel, seq_len)
        """
        out = self.network(x)
        # return self.network(x).T
        return out


class TCN(nn.Module):
    def __init__(self, input_size, output_size, channels, feature_dim, kernel_size, stride, dropout):
        """
        TCN面对外部的包装类

        :param input_size: int, 输入通道数
        :param output_size: int, 输出步长
        :param channels: list, 每层的hidden_channel数，例如[25,25,25,25]表示有4个隐层，每层hidden_channel数为25
        :param kernel_size: int, 卷积核尺寸
        :param stride: int, 卷积步长
        :param dropout: float, drop_out比率
        """
        super(TCN, self).__init__()
        self.tcn = TemporalConvNet(input_size, channels, kernel_size, stride, dropout)
        self.fc = nn.Sequential(
            nn.Flatten(0, -1),
            nn.Linear(channels[-1] * feature_dim, output_size)
        )

    def forward(self, x):
        out = self.tcn(x)
        out = self.fc(out).reshape(1, -1)
        return out


'''
def test(dataloader_train):
    model = TCN()

    data = next(iter(dataloader_train))
    # 应GCN的数据类型要求,将 float64(Float) -> float32(Double)
    data = data.to(torch.float32)
    # input与label
    input = data[0:(batch_size - output_size), :]
    label = data[(batch_size - output_size):, 0].view(1, output_size)
    print(f"input: {input.shape}")
    print(f"label: {label.shape}")

    out = model(input)
    print(f"out: {out.shape}")
'''









